--[[## ft=lua ##################################################################
  This file is licensed under [WTFPL v2](https://www.wtfpl.net/).  You may read
  the full license in the file LICENSE or by following the web link.

  TABLEX - extensions to table operations in Lua.

  This is a collection of miscellaneous routines to use as a drop-in replacement
  for the standard table library.  For convenience, the standard table library
  is shadowed thus permitting the following in code:

    local table = require 'tablex'

  This lets all current table operations get used in parallel with the
  extensions in a seamless fashion.
  ############################################################################]]

--[[####################################
  Shadow the standard library.
--####################################]]
local tbl = table
local mt = {
  __index = tbl
}
local M = setmetatable({}, mt)

--[[####################################
  Helpful local utilities
--####################################]]

--[[################
  Define a total ordering for keys:
  - numbers first
  - then strings
  - then all other key types, ordered as if strings
--################]]
local function order(k1, k2)
    local type_k1 = type(k1)
    local type_k2 = type(k2)
    if (type_k1 == 'string' and type_k2 == 'string') or
       (type_k1 == 'number' and type_k2 == 'number') then
        return k1 < k2
    elseif type_k1 == 'number' and type_k2 ~= 'number' then
        return true
    elseif type_k1 == 'string' and type_k2 ~= 'number' then
        return false
    else
        return tostring(k1) < tostring(k2)
    end
end

--[[####################################
  Table shuffling API.
--####################################]]

--[[################
  In-place array shuffle.  shufflei(a) takes an array-flavoured table, shuffles
  it in-place, then returns that table (for convenient chaining, nesting, etc.).
--################]]
local function shufflei(a)
  assert(type(a) == 'table', 'can only shuffle tables')
  local i
  for i = 1, #a do
    local s = math.random(i, #a)
    a[i],a[s] = a[s],a[i]
  end
  return a
end

--[[################
  Shuffle.  shuffle(a) takes an array-flavoured table, copies it, shuffles the
  copy and returns it.
--################]]
local function shuffle(a)
  assert(type(a) == 'table', 'can only shuffle tables')
  local r = {}
  for _,v in ipairs(a) do tbl.insert(r, v) end
  return shufflei(r)
end

--[[################
  Shuffle in-place.  kvshufflei(t) takes a table of K-V pairs (a general table)
  and mixes up the values such that the former K-V relationships are mixed up,
  analogous to how they'd be mixed up in a numerically-indexed table's shuffling.
--################]]
local function kvshufflei(t)
  assert(type(t) == 'table', 'can only shuffle tables')
  local ks = {}
  local vs = {}
  for k,v in pairs(t) do
    tbl.insert(ks, k)
    tbl.insert(vs, v)
  end
  shufflei(ks)
  local i
  for i = 1, #ks do
    t[ks[i]] = vs[i]
  end
  return t
end

--[[################
  Shuffle.  kvshuffle(t) takes a table of K-V pairs (a general table)  and
  mixes up the values such that the former K-V relationships are mixed up,
  analogous to how they'd be mixed up in a numerically-indexed table's shuffling.
--################]]
local function kvshuffle(t)
  assert(type(t) == 'table', 'can only shuffle tables')
  local r = {}
  for k,v in pairs(t) do r[k] = v end
  return kvshufflei(r)
end

M.kvshuffle = kvshuffle
M.kvshufflei = kvshufflei
M.shuffle = shuffle
M.shufflei = shufflei

--[[####################################
  Sorted-order K-V table iterator.
  - Can cope with key/value deletion mid-traversal.
  - Can cope with key/value addition mid-traversal but will not traverse added
    keys.
  - CANNOT cope with keys that are not sortable in table.sort() unless a custom
    comparator is provided.
--####################################]]

--[[################
  Iterator.
--################]]
local function spairs(t, c)

    local function make_snext()
        local function sindex()
            local skeys = {}

            for k in pairs(t) do
                tbl.insert(skeys, k)
            end

            tbl.sort(skeys, c)

            local i, j = 1, #skeys

            while i < j do
                skeys[i], skeys[j] = skeys[j], skeys[i]
                i = i + 1
                j = j - 1
            end

            return skeys
        end

        local states = setmetatable(sindex(), {__mode = 'k'})

        local function snext(s)
            local k = tbl.remove(states)

            if k == nil then return end

            local v = rawget(t, k)
            if v ~= nil then
                return k, v
            else
                return snext(s)
            end
        end

        return snext, t, nil
    end

    return make_snext()
end

M.spairs = spairs

--[[####################################
  General table utilities.
--####################################]]

--[[################
  Dump a table's contents (including metatables) recursively.
--################]]
local function dump(t, n, i)
    assert(type(t) == 'table', 'can only dump tables')

    local indent = i or 0
    local function dump_value(k, v)
        return string.rep(' ', indent) .. tostring(k) .. ' -> ' .. tostring(v) .. '\n'
    end

    local function dump_table(n, t)
        local name = n or '(unnamed)'
        local rv = ''
        rv = rv .. string.rep(' ', indent) .. tostring(name) .. ' -> ' .. tostring(t) .. '\n'
        indent = indent + 2
        for k,v in spairs(t, order) do
--        for k,v in pairs(t) do
            if type(v) == 'table' then
                rv = rv .. dump(v, k, indent)
            else
                rv = rv .. dump_value(k, v)
            end
        end
        return rv
    end

    return dump_table(n, t)
end

--[[################
  Execute a function for each element of a table.
  This should have been in the higher-order section, but functions here need it
  declared in advance without the clumsiness of forward declaration
--################]]
local function foreach(t, f, ...)
  for k,v in pairs(t) do
    f(k, v, ...)
  end
end

--[[################
  Count the number of times a given element appears in a table.
--################]]
local function count(t, e)
  local r = 0
  foreach(t, function(_,v) if v == e then r = r + 1 end end)
  return r
end

--[[################
  Find the first key (searched in arbitrary order) whose value matches an
  element.
--################]]
local function find(t, e)
  for k,v in pairs(t) do
    if v == e then
      return k
    end
  end
end

--[[################
  Return the number of entries in a table.
--################]]
local function size(t)
  local r = 0
  for _,_ in pairs(t) do
    r = r + 1
  end
  return r
end

M.dump = dump
M.count = count
M.find = find
M.size = size

--[[####################################
  General array utilities.
--####################################]]

--[[################
  Slice an array-flavoured table, left and right.
--################]]
local function slice(t, l, r)
  if l < 0      then error 'cannot skip negative numbers of arguments' end
  if r < 0      then error 'cannot drop negative numbers of arguments' end
  if l + r > #t then error 'cannot slice more elements than arguments' end
  return { tbl.unpack(t, l+1, #t-r) }
end

--[[################
  Skip the first elements of an array-flavoured table.
--################]]
local function skip(t, l)
  return slice(t, l, 0)
end

--[[################
  Drop the last elements of an array-flavoured table.
--################]]
local function drop(t, r)
  return slice(t, 0, r)
end

M.slice = slice
M.skip = skip
M.drop = drop

--[[####################################
  Higher-order operations.
--####################################]]

-- note that foreach belongs here, but we use it earlier and local predeclaration syntax is clunky

--[[################
  Map a function to a table's values.
--################]]
local function map(t, f)
  assert(type(t) == 'table', 'can only map over tables')
  local r = {}
  foreach(t, function(k,v)
    r[k] = f(v)
  end)
  return r
end

--[[################
  Filter values by a user-provided predicate.
--################]]
local function filter(t, p, ...)
  assert(type(t) == 'table', 'can only filter tables')
  local r = {}
  foreach(t, function(k, v, ...)
    if p(k, v, ...) then
        r[k] = v
    end
  end)
  return r
end

M.map = map
M.filter = filter
M.foreach = foreach

return setmetatable({}, { __index = M })
